package cn.quang.utils.lang;

import java.util.*;
import java.util.function.Function;

public class CollectionUtils {

    /**
     * 传入集合A和集合B，使集合B的元素完全按照集合A的顺序排列。
     * @param reference 基准List集合
     * @param dataList 需要排序的List集合
     * @param function 用于将集合B的元素映射到集合A的元素
     * @return 排序后的List集合
     */
    public static <T, K> List<T> sortByReference(List<K> reference, List<T> dataList, Function<T, K> function) {
        if (reference == null || reference.isEmpty() || dataList == null || dataList.isEmpty() ) {
            return dataList;
        }
        int maxSize = reference.size() + dataList.size() + 1;
        Map<K, Integer> sortIndexMap = new HashMap<>();
        for (K sortField : reference) {
            sortIndexMap.put(sortField, maxSize--);
        }
        Map<Integer, T> resultMap = new TreeMap<>(Comparator.reverseOrder());
        for (T data : dataList) {
            K fieldName = function.apply(data);
            int index = sortIndexMap.getOrDefault(fieldName, 0);
            if (index > 0) {
                resultMap.put(index, data);
            } else {
                resultMap.put(maxSize--, data);
            }
        }
        return new ArrayList<>(resultMap.values());
    }
}
